# voice_input.py
import whisper
import sounddevice as sd
import wavio
import sys
import os
import pyperclip
import time

# 配置参数
MODEL_NAME = "medium" # 我们选择的 Whisper 模型
SAMPLE_RATE = 16000  # 语音识别推荐的采样率
CHANNELS = 1        # 单声道
TEMP_AUDIO_FILE = "D:\\Py\\python\\Script\\temp_recording.wav" # 临时音频文件路径

def record_audio(duration):
    """录制指定时长的音频"""
    print(f"开始录音 {duration} 秒...")
    # 确保录音设备可用
    try:
        sd.query_devices(kind='input')
    except Exception as e:
        print(f"错误: 无法找到音频输入设备或设备不可用。请检查麦克风设置。{e}", file=sys.stderr)
        return None

    audio_data = sd.rec(int(duration * SAMPLE_RATE), samplerate=SAMPLE_RATE, channels=CHANNELS, dtype='int16')
    sd.wait() # 等待录音完成
    wavio.write(TEMP_AUDIO_FILE, audio_data, SAMPLE_RATE, sampwidth=2)
    print(f"录音完成，保存到 {TEMP_AUDIO_FILE}")
    return TEMP_AUDIO_FILE

def transcribe_and_copy(audio_file):
    """使用 Whisper 识别音频并复制到剪贴板"""
    if not os.path.exists(audio_file):
        print(f"错误: 音频文件 '{audio_file}' 不存在。", file=sys.stderr)
        return

    try:
        print(f"加载 Whisper 模型: {MODEL_NAME}...")
        model = whisper.load_model(MODEL_NAME) # 这里会检查并下载模型
        
        print("开始识别...")
        
        # 核心修改：在 transcribe 方法中加入 initial_prompt 参数，引导输出简体中文
        result = model.transcribe(
            audio_file, 
            language="zh",                  # 明确指定语言为中文
            initial_prompt="这是一段普通话，请输出简体中文。", # 引导模型输出简体
            fp16=False                      # 如果没有GPU或遇到问题，可以设置为False，确保CPU兼容性
        )
        
        transcribed_text = result["text"].strip() # 移除首尾空格
        
        print(f"识别结果: {transcribed_text}")
        
        if transcribed_text:
            pyperclip.copy(transcribed_text)
            print("识别结果已复制到剪贴板。")
        else:
            print("未识别到有效文本。")

    except Exception as e:
        print(f"发生识别或模型加载错误: {e}", file=sys.stderr)
    finally:
        # 清理临时文件
        if os.path.exists(audio_file):
            os.remove(audio_file)

if __name__ == "__main__":
    # 允许通过命令行参数指定录音时长，默认 5 秒
    record_duration = float(sys.argv[1]) if len(sys.argv) > 1 else 10.0
    
    recorded_file = record_audio(record_duration)
    if recorded_file:
        transcribe_and_copy(recorded_file)